package com.yupi.springbootinit.spark;

import cn.hutool.json.JSONUtil;
import com.yupi.springbootinit.model.dto.spark.ChatRequest;
import com.yupi.springbootinit.model.vo.ChatResponse;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import okhttp3.Response;
import okhttp3.WebSocket;
import okhttp3.WebSocketListener;

import java.io.IOException;
import java.util.Collections;
import java.util.List;

@Slf4j
@Data
public class SparkChatListener extends WebSocketListener {
    public StringBuilder totalAnswer = new StringBuilder();

    /**
     * websocket连接关闭标志
     */
    private Boolean wsCloseFlag;
    @Override
    public void onOpen(WebSocket webSocket, Response response) {
        super.onOpen(webSocket, response);
        Thread thread = new Thread(new Runnable() {
            @Override
            public void run() {
                try {
                    ChatRequest chatRequest = new ChatRequest();
                    ChatRequest.builder()
                            .header(ChatRequest.Header.builder()
//                                    .appId(appid)
//                                    .uid((int) userId)
                                    .build())
                            .parameter(ChatRequest.Parameter.builder()
                                    .chat(ChatRequest.Chat.builder()
//                                            .domain(domain)
                                            .temperature(0.5)
                                            .maxTokens(4096)
                                            .build())
                                    .build())
                            .payload(ChatRequest.Payload.builder()
                                    .message(ChatRequest.Message
                                            .builder()
                                            .text(Collections.singletonList(
                                                    ChatRequest.Text.builder()
//                                                            .content(NewQuestion)
                                                            .role("user")
                                                            .build()))
                                            .build())
                                    .build());
                    webSocket.send(chatRequest.toString());
                    // 等待服务端返回完毕后关闭
                    while (true) {
                        // System.err.println(wsCloseFlag + "---");
                        Thread.sleep(200);
                        if (wsCloseFlag) {
                            break;
                        }
                    }
                    webSocket.close(1000, "");
//                   thread.interrupt();
                } catch (Exception e) {
                    e.printStackTrace();
                }
            }
        });
        thread.start();
    }

    @Override
    public void onMessage(WebSocket webSocket, String text) {
        // System.out.println(userId + "用来区分那个用户的结果" + text);
        ChatResponse chatResponse = JSONUtil.toBean(text, ChatResponse.class);
        if (chatResponse.getHeader().getCode() != 0) {
            System.out.println("发生错误，错误码为：" + chatResponse.getHeader().getCode());
            System.out.println("本次请求的sid为：" + chatResponse.getHeader().getSid());
            webSocket.close(1000, "");
        }
        List<ChatResponse.ContentRoleIndex> textList = chatResponse.getPayload().getChoices().getText();
        for (ChatResponse.ContentRoleIndex temp : textList) {
            totalAnswer.append(temp.getContent());
        }
        if (chatResponse.getHeader().getStatus() == 2) {
            // 可以关闭连接，释放资源
            ChatResponse.ContentRoleIndex contentRoleIndex = new ChatResponse.ContentRoleIndex();
            contentRoleIndex.setRole("assistant");
            contentRoleIndex.setContent(totalAnswer.toString());
            wsCloseFlag = true;
        }
    }

    @Override
    public void onFailure(WebSocket webSocket, Throwable t, Response response) {
        super.onFailure(webSocket, t, response);
        if (response != null) {
            int code = response.code();
            try {
                log.info("[Spark-chat-websocket]failure code:{} , boyd{}", code, response.body().string());
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        }
    }
}